package util

//@author: luodesheng 308189510@qq.com
//@create: 2024-09-03 by vscode

import (
	"bytes"
	"crypto"
	"crypto/aes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
)

func AesEncrypt(key []byte, data []byte) []byte {
	b, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}
	blockSize := b.BlockSize()
	data = PKCS7Padding(data, blockSize)

	// cbc模式
	//bm := cipher.NewCBCEncrypter(b, key[:blockSize])

	r := make([]byte, len(data))
	tmp := r
	for len(data) > 0 {
		b.Encrypt(tmp, data[:blockSize])
		data = data[blockSize:]
		tmp = tmp[blockSize:]
	}
	//bm.CryptBlocks(r, data)
	return r
}

func AesDecrypt(key []byte, data []byte) []byte {
	b, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}
	blockSize := b.BlockSize()

	//	bm := cipher.NewCBCDecrypter(b, key[:blockSize])

	r := make([]byte, len(data))
	tmp := r
	for len(data) > 0 {
		b.Decrypt(tmp, data[:blockSize])
		data = data[blockSize:]
		tmp = tmp[blockSize:]
	}

	//	bm.CryptBlocks(r, data)

	return PKCS7UnPadding(r)
}

func PKCS7Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padtext...)
}

func PKCS7UnPadding(origData []byte) []byte {
	length := len(origData)
	unpadding := int(origData[length-1])
	return origData[:(length - unpadding)]
}

func RsaEncrypt(publicKey *rsa.PublicKey, data []byte) []byte {
	b, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, data)
	if err != nil {
		panic(err)
	}
	return b
}

func GetPublicKey(publicKey []byte) *rsa.PublicKey {
	pub, err := x509.ParsePKIXPublicKey(publicKey)
	if err != nil {
		panic(err)
	}
	return pub.(*rsa.PublicKey)
}

// pkcs8 格式私钥
func GetPrivateKey(privateKey []byte) *rsa.PrivateKey {
	pri, err := x509.ParsePKCS8PrivateKey(privateKey)
	if err != nil {
		panic(err)
	}
	return pri.(*rsa.PrivateKey)
}

func RsaDecrypt(privateKey *rsa.PrivateKey, data []byte) []byte {
	b, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, data)
	if err != nil {
		panic(err)
	}
	return b
}

func Sha256(src []byte) []byte {
	h := sha256.New()
	h.Write(src)
	return h.Sum(nil)
}

func VerifyWithSha256(publicKey *rsa.PublicKey, data []byte, sign []byte) bool {
	err := rsa.VerifyPKCS1v15(publicKey, crypto.SHA256, Sha256(data), sign)
	return err == nil
}

func SignWithSha256(privateKey *rsa.PrivateKey, data []byte) []byte {
	b, err := rsa.SignPKCS1v15(rand.Reader, privateKey, crypto.SHA256, Sha256(data))
	if err != nil {
		panic(err)
	}
	return b
}
